import os

from langchain_core.runnables import RunnableConfig

from owl_ai.application import ImageUtils
from owl_ai.domain.graph.nodes.graph_node import BaseGraphNode
from owl_ai.mapper import ChatFileMapper
from owl_common.config import OWLConfig


class DocumentExtractNode(BaseGraphNode):
    """
    文件解析节点，提供文档内容抽取功能，文本文件抽取为文本，图像抽取为base64，返回的数据格式为：
    {
        "type": "文件类型: txt: 文本类， image：图像类",
        "suffix": "文件后缀",
        "name": "文件名",
        "data": "内容"
    }
    """

    def __init__(self, node_config: dict):
        self.node_id = node_config.get("id")
        self.inputs = node_config.get("inputs")
        self.outputs = node_config.get("outputs")

    @classmethod
    def file_extract(cls, chat_file):
        """
        文件解析
        Args:
            file_path: 文件路径
        Returns:
            解析后的内容
        """
        # 实现文件解析逻辑
        file_name = chat_file.name
        file_suffix = file_name.split(".")[-1]
        file_type = ''
        file_data = ''

        if ImageUtils.is_image_file(file_name):
            file_type = 'image'
            file_path = os.path.join(OWLConfig().upload_path, chat_file.path)
            file_data = ImageUtils.convert_to_base64(file_path)

        return file_type, file_suffix, file_name, file_data

    def __call__(self, state, config: RunnableConfig):
        input_params = self.input_params_assem(state, self.inputs)
        app = state['app']

        parse_files = {}
        for param_key in input_params.keys():
            param_values = input_params.get(param_key)

            item_files = param_values.get('files')
            if item_files:
                params_files = []
                for file_id in item_files:

                    with app.app_context():
                        chat_file = ChatFileMapper.find_file_by_id(file_id)

                    file_type, file_suffix, file_name, file_data = self.file_extract(chat_file)

                    params_files.append({
                        'type': file_type,
                        'suffix': file_suffix,
                        'name': file_name,
                        'data': file_data
                    })
                if len(params_files) > 0:
                    parse_files.setdefault(param_key, params_files)

        return {
            "node_params": {
                self.node_id: {
                    "files": parse_files
                }
            }
        }
